import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import DataLoader, TensorDataset

torch.manual_seed(42)


# 1. 生成随机数据
def generate_data(num_samples=1000, input_size=10):
    """生成随机输入数据和二分类标签"""
    X = torch.randn(num_samples, input_size)  # 输入数据 (1000x10)
    Y = torch.randint(0, 2, (num_samples, 1)).float()  # 二分类标签 (0或1)
    return X, Y


# 2. 定义网络结构
class DynamicNet(nn.Module):
    def __init__(self, input_size, hidden_size, output_size):
        super(DynamicNet, self).__init__()
        self.fc1 = nn.Linear(input_size, hidden_size)  # 输入层 → 隐藏层
        self.relu = nn.ReLU()  # ReLU激活函数
        self.fc2 = nn.Linear(hidden_size, output_size)  # 隐藏层 → 输出层
        self.sigmoid = nn.Sigmoid()  # Sigmoid激活函数（用于二分类）

    def forward(self, x):
        x = self.fc1(x)
        x = self.relu(x)
        x = self.fc2(x)
        x = self.sigmoid(x)
        return x


# 3. 训练函数
def train(model, train_loader, criterion, optimizer, epochs=100):
    model.train()  # 设置为训练模式
    for epoch in range(epochs):
        print(f"第 {epoch + 1} 轮训练中...")
        train_loss = 0.0
        for inputs, labels in train_loader:
            optimizer.zero_grad()  # 清零梯度
            outputs = model(inputs)  # 前向传播
            loss = criterion(outputs, labels)  # 计算损失
            loss.backward()  # 反向传播
            optimizer.step()  # 更新参数
            train_loss += loss.item()
        print(f"第 {epoch + 1} 轮训练结束，平均损失：{train_loss / len(train_loader):.4f}")


# 4. 测试函数
def model_test(model, test_loader):
    model.eval()  # 设置为评估模式
    correct = 0
    total = 0
    with torch.no_grad():  # 不计算梯度
        for inputs, labels in test_loader:
            outputs = model(inputs)
            predicted = (outputs > 0.5).float()  # 预测类别（0或1）
            total += labels.size(0)
            correct += (predicted == labels).sum().item()
    accuracy = 100 * correct / total
    print(f"测试集准确率：{accuracy:.2f}%")


# 5. 主函数
def main01():
    print("=======main01=======")
    # 生成数据
    X_train, Y_train = generate_data(num_samples=800, input_size=10)
    X_test, Y_test = generate_data(num_samples=200, input_size=10)

    # 创建 DataLoader
    train_dataset = TensorDataset(X_train, Y_train)
    test_dataset = TensorDataset(X_test, Y_test)
    train_loader = DataLoader(train_dataset, batch_size=32, shuffle=True)
    test_loader = DataLoader(test_dataset, batch_size=32, shuffle=False)

    # 初始化模型、损失函数和优化器
    model = DynamicNet(input_size=10, hidden_size=64, output_size=1)
    criterion = nn.BCELoss()  # 二分类交叉熵损失  测试集准确率：51.50%
    # criterion = nn.MSELoss()  # 二分类交叉熵损失 测试集准确率：51.50%
    optimizer = optim.Adam(model.parameters(), lr=0.001)

    # 训练和测试
    train(model, train_loader, criterion, optimizer, epochs=5)
    model_test(model, test_loader)


if __name__ == "__main__":
    main01()
